import mindspore.nn as nn
from mindspore.ops import operations as P




class ResidualUnit(nn.Cell):
    def __init__(self, in_channel, out_channel, stride=2, kernel_size=(3, 3, 3), down=True, is_output=False):
        super().__init__()
        self.stride = stride
        self.down = down
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.down_conv_1 = nn.Conv3d(in_channel, out_channel, kernel_size=(3, 3, 3), \
                                     pad_mode="pad", stride=self.stride, padding=1)
        self.is_output = is_output
        if not is_output:
            self.batchNormal1 = nn.BatchNorm3d(num_features=self.out_channel)
            self.relu1 = nn.PReLU()
        if self.down:
            self.down_conv_2 = nn.Conv3d(out_channel, out_channel, kernel_size=(3, 3, 3), \
                                         pad_mode="pad", stride=1, padding=1)
            self.relu2 = nn.PReLU()
            if kernel_size[0] == 1:
                self.residual = nn.Conv3d(in_channel, out_channel, kernel_size=(1, 1, 1), \
                                          pad_mode="valid", stride=self.stride)
            else:
                self.residual = nn.Conv3d(in_channel, out_channel, kernel_size=(3, 3, 3), \
                                          pad_mode="pad", stride=self.stride, padding=1)
            self.batchNormal2 = nn.BatchNorm3d(num_features=self.out_channel)


    def construct(self, x):
        out = self.down_conv_1(x)
        if self.is_output:
            return out
        out = self.batchNormal1(out)
        out = self.relu1(out)
        if self.down:
            out = self.down_conv_2(out)
            out = self.batchNormal2(out)
            out = self.relu2(out)
            res = self.residual(x)
        else:
            res = x
        return out + res

class Down(nn.Cell):
    def __init__(self, in_channel, out_channel, stride=2, kernel_size=(3, 3, 3)):
        super().__init__()
        self.stride = stride
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.down_conv = ResidualUnit(self.in_channel, self.out_channel, stride, kernel_size)

    def construct(self, x):
        x = self.down_conv(x)
        return x


class Up(nn.Cell):
    def __init__(self, in_channel, down_in_channel, out_channel, stride=2, is_output=False):
        super().__init__()
        self.in_channel = in_channel
        self.down_in_channel = down_in_channel
        self.out_channel = out_channel
        self.stride = stride
        self.conv3d_transpose = nn.Conv3dTranspose(in_channels=self.in_channel + self.down_in_channel, \
                                                   out_channels=self.out_channel, kernel_size=(3, 3, 3), \
                                                   pad_mode="pad", stride=self.stride, \
                                                   output_padding=(1, 1, 1), padding=1)

        self.concat = P.Concat(axis=1)
        self.conv = ResidualUnit(self.out_channel, self.out_channel, stride=1, down=False, \
                                 is_output=is_output)
        self.batchNormal1 = nn.BatchNorm3d(num_features=self.out_channel)
        self.relu = nn.PReLU()

    def construct(self, input_data, down_input):
        x = self.concat((input_data, down_input))
        x = self.conv3d_transpose(x)
        x = self.batchNormal1(x)
        x = self.relu(x)
        x = self.conv(x)
        return x

class UNet3d(nn.Cell):
    def __init__(self, in_channels=1, base_channels=16, num_classes=2):
        super(UNet3d, self).__init__()
        self.n_channels = in_channels
        self.n_classes = num_classes

        # down
        self.down1 = Down(in_channel=self.n_channels, out_channel=base_channels)
        self.down2 = Down(in_channel=base_channels, out_channel=base_channels*2)
        self.down3 = Down(in_channel=base_channels*2, out_channel=base_channels*4)
        self.down4 = Down(in_channel=base_channels*4, out_channel=base_channels*8)
        self.down5 = Down(in_channel=base_channels*8, out_channel=base_channels*16, stride=1, kernel_size=(1, 1, 1))

        # up
        self.up1 = Up(in_channel=base_channels*16, down_in_channel=base_channels*8, out_channel=base_channels*4)
        self.up2 = Up(in_channel=base_channels*4, down_in_channel=base_channels*4, out_channel=base_channels*2)
        self.up3 = Up(in_channel=base_channels*2, down_in_channel=base_channels*2, out_channel=base_channels)
        self.up4 = Up(in_channel=base_channels, down_in_channel=base_channels, out_channel=self.n_classes, is_output=True)


    def construct(self, input_data):
        x1 = self.down1(input_data)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.down5(x4)

        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        return x
